// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <turbo/log/logging.h>
#include <gtest/gtest.h>
#include <cstdint>

#include <pollux/expression/vector_readers.h>
#include <pollux/testing/functions/prestosql/function_base_test.h>
#include <pollux/vector/base_vector.h>

namespace kumo::pollux::exec {

class ConstantFlatVectorReaderTest : public functions::test::FunctionBaseTest {
 public:
  ConstantFlatVectorReader<int32_t> makeConstantFlatVectorReader(
      const VectorPtr& vector) {
    auto reader = vector->encoding() == VectorEncoding::Simple::FLAT
        ? ConstantFlatVectorReader<int32_t>(
              static_cast<FlatVector<int32_t>*>(vector.get()))
        : ConstantFlatVectorReader<int32_t>(
              static_cast<ConstantVector<int32_t>*>(vector.get()));
    return reader;
  }
};

template <typename T>
void testFlatContainsNoNulls(const T& reader, const VectorPtr& vector) {
  ASSERT_FALSE(reader.may_have_nulls());
  ASSERT_FALSE(reader.may_have_nulls_recursive());

  for (vector_size_t i = 0; i < vector->size(); i++) {
    ASSERT_EQ(i * 2, reader[i]);
    ASSERT_EQ(i * 2, reader.readNullFree(i));
    ASSERT_TRUE(reader.isSet(i));
    ASSERT_FALSE(reader.containsNull(i));
    ASSERT_FALSE(reader.containsNull(0, i));
    ASSERT_FALSE(reader.containsNull(i, vector->size()));
  }
}

TEST_F(ConstantFlatVectorReaderTest, flatContainsNoNulls) {
  auto vector =
      make_flat_vector<int32_t>(10, [](vector_size_t row) { return row * 2; });

  FlatVectorReader<int32_t> reader1(*vector);
  auto reader2 = this->makeConstantFlatVectorReader(vector);

  testFlatContainsNoNulls(reader1, vector);
  testFlatContainsNoNulls(reader2, vector);
}

template <typename T>
void testFlatContainsNull(const T& reader, const VectorPtr& vector) {
  ASSERT_TRUE(reader.may_have_nulls());
  ASSERT_TRUE(reader.may_have_nulls_recursive());

  for (vector_size_t i = 0; i < vector->size(); i++) {
    if (i % 5 == 2) {
      ASSERT_FALSE(reader.isSet(i));
      ASSERT_TRUE(reader.containsNull(i));
    } else {
      ASSERT_EQ(i * 2, reader[i]);
      ASSERT_EQ(i * 2, reader.readNullFree(i));
      ASSERT_TRUE(reader.isSet(i));
      ASSERT_FALSE(reader.containsNull(i));
    }

    if (i > 2) {
      ASSERT_TRUE(reader.containsNull(0, i));
    } else {
      ASSERT_FALSE(reader.containsNull(0, i));
    }

    if (i <= 7) {
      ASSERT_TRUE(reader.containsNull(i, vector->size()));
    } else {
      ASSERT_FALSE(reader.containsNull(i, vector->size()));
    }
  }
}

TEST_F(ConstantFlatVectorReaderTest, flatContainsNulls) {
  auto vector = make_flat_vector<int32_t>(
      10,
      [](vector_size_t row) { return row * 2; },
      [](vector_size_t row) { return row % 5 == 2; });

  FlatVectorReader<int32_t> reader1(*vector);
  auto reader2 = this->makeConstantFlatVectorReader(vector);
  testFlatContainsNull(reader1, vector);
  testFlatContainsNull(reader2, vector);
}

template <typename T>
void testConstant(const T& reader, const VectorPtr& vector) {
  ASSERT_FALSE(reader.may_have_nulls());
  ASSERT_FALSE(reader.may_have_nulls_recursive());

  for (vector_size_t i = 0; i < vector->size(); i++) {
    ASSERT_EQ(5, reader[i]);
    ASSERT_EQ(5, reader.readNullFree(i));
    ASSERT_TRUE(reader.isSet(i));
    ASSERT_FALSE(reader.containsNull(i));
    ASSERT_FALSE(reader.containsNull(0, i));
    ASSERT_FALSE(reader.containsNull(i, vector->size()));
  }
}

TEST_F(ConstantFlatVectorReaderTest, constant) {
  auto vector = make_constant<int32_t>(5, 10);

  ConstantVectorReader<int32_t> reader1(*vector->as<ConstantVector<int32_t>>());
  auto reader2 = this->makeConstantFlatVectorReader(vector);
  testConstant(reader1, vector);
  testConstant(reader2, vector);
}

template <typename T>
void testConstantNull(const T& reader, const VectorPtr& vector) {
  ASSERT_TRUE(reader.may_have_nulls());
  ASSERT_TRUE(reader.may_have_nulls_recursive());

  for (vector_size_t i = 0; i < vector->size(); i++) {
    ASSERT_FALSE(reader.isSet(i));
    ASSERT_TRUE(reader.containsNull(i));
    if (i > 0) {
      ASSERT_TRUE(reader.containsNull(0, i));
    }
    ASSERT_TRUE(reader.containsNull(i, vector->size()));
  }
}

TEST_F(ConstantFlatVectorReaderTest, constantNull) {
  auto vector = make_constant<int32_t>(std::nullopt, 10);

  ConstantVectorReader<int32_t> reader1(*vector->as<ConstantVector<int32_t>>());
  auto reader2 = this->makeConstantFlatVectorReader(vector);

  testConstantNull(reader1, vector);
  testConstantNull(reader2, vector);
}
} // namespace kumo::pollux::exec
